import argparse
from argparse import Namespace
import os

from eval_tool.third_party.hloc.hloc.localize_inloc import localize_with_matcher
from immatch.utils.model_helper import init_model

def eval_inloc(args):
    model, model_conf = init_model(args.config, args.benchmark_name)
    matcher = lambda im1, im2: model.match_pairs(im1, im2)
    args = Namespace(**vars(args), **model_conf)
    print(">>>>", args)

    # Setup output dir
    rthres = args.rthres
    mthres = args.match_threshold
    skip_matches = args.skip_matches
    retrieval_pairs = os.path.join(args.pair_dir, args.benchmark_name, args.pairs)
    pair_tag = retrieval_pairs.split('query-')[-1].replace('.txt', '')
    exp_name = f'{pair_tag}_sk{skip_matches}im{args.imsize}rth{rthres}mth{mthres}'
    if args.prefix:
        exp_name = f'{exp_name}.{args.prefix}'
    odir = os.path.join('outputs', args.benchmark_name, model.name, exp_name)
    method_tag = f'{model.name}_{exp_name}'        
    print('>>>Method tag:', method_tag)

    # Localize InLoc queries
    localize_with_matcher(
        matcher, args.dataset_dir, retrieval_pairs, odir,
        method_tag, rthres=rthres, skip_matches=skip_matches
    )
    
if __name__ == '__main__':    
    parser = argparse.ArgumentParser(description='Localize Inloc')
    parser.add_argument('--gpu', '-gpu', type=str, default='0')
    parser.add_argument('--loftr_config', type=str, default=None)
    parser.add_argument('--prefix', type=str, default=None)
    parser.add_argument('--dataset_dir', type=str, default='data/datasets/InLoc')
    parser.add_argument('--pair_dir', type=str, default='data/pairs')
    parser.add_argument('--benchmark_name', type=str, default='inloc')

    args = parser.parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    eval_inloc(args)
    